import threading
import gevent
import time
import logging

_WARNING = logging.getLogger(__name__).warning

class TaskThread(threading.Thread):
    IDLE = 1
    BUSY = 2

    def __init__(self):
        super(TaskThread, self).__init__()
        self.task_queue = []
        self.force_exit = False
        self.lock = threading.Lock()
        self.work_signal = threading.Semaphore(value=0)
        self.work_status = self.IDLE

    def add_task(self, call_func, args):
        self.lock.acquire()
        self.task_queue.append((call_func, args))
        self.lock.release()

    def task_count(self):
        return len(self.task_queue)

    def run(self):
        while not self.force_exit or self.task_queue:
            if self.task_queue:
                spawns = []
                task_count = len(self.task_queue)
                for i in range(task_count):
                    f, p = self.task_queue[i]
                    spawns.append(gevent.spawn(f, *p))
                gevent.joinall(spawns)

                self.lock.acquire()
                del self.task_queue[0:task_count]
                self.lock.release()
            else:
                # 进入空闲状态
                self.work_status = self.IDLE
                self.work_signal.acquire()  # 等待下次任务的唤醒

        print("thread: {} exit".format(threading.currentThread().ident))

    def shutdown(self):
        self.force_exit = True


class TaskThreadPool:
    def __init__(self, thread_count):
        self.stoped = False
        self.task_threads = []
        for i in range(thread_count):
            task_thread = TaskThread()
            task_thread.start()
            self.task_threads.append(task_thread)

    def get_work_thread(self):
        """得到一个工作线程"""
        for task_thread in self.task_threads:
            if task_thread.work_status == TaskThread.IDLE:
                task_thread.work_status = TaskThread.BUSY
                return task_thread

        task_index = 0
        total_task_count = 0
        min_task_count = self.task_threads[0].task_count()
        for i in range(1, len(self.task_threads)):
            tc = self.task_threads[i].task_count()
            total_task_count += tc
            if tc < min_task_count:
                task_index = i
                min_task_count = tc

        if (total_task_count / len(self.task_threads)) > 3:
            _WARNING("task to many, count is: {}".format(total_task_count))

        return self.task_threads[task_index]

    def add_task(self, call_func, *args):
        task_thread = self.get_work_thread()
        task_thread.add_task(call_func, args)
        task_thread.work_signal.release()  # 触发执行

    def add_tasks(self, tasks):
        task_thread = self.get_work_thread()
        for cf, args in tasks:
            task_thread.add_task(cf, args)
            task_thread.work_signal.release()  # 触发执行

    def shutdown(self):
        if self.stoped:
            return

        self.stoped = True

        for task_thread in self.task_threads:
            task_thread.shutdown()
            task_thread.work_signal.release()  # 触发执行，让线程退出

        for task_thread in self.task_threads:
            task_thread.join()

    def sleep(self, seconds):
        gevent.sleep(seconds)

    def task_status(self):
        if self.stoped:
            return None

        result = []
        for task_thread in self.task_threads:
            if task_thread.work_status == TaskThread.IDLE:
                result.append((task_thread.ident, TaskThread.IDLE, "idle"))
            else:
                result.append((task_thread.ident, TaskThread.BUSY, "busy"))
        return result


TASK_QUEUE = TaskThreadPool(4)


if __name__ == "__main__":
    def test_func1(task_id, b, c):
        print("task_id = {}, b = {}, c = {}, thread_id = {}".format(
            task_id, b, c, threading.currentThread().ident))
        TASK_QUEUE.sleep(1)
        print("task : {} sleep end exit".format(task_id))

    for i in range(20):
        TASK_QUEUE.add_task(test_func1, i, 1, 2)

    print("task dispatch stop")
    time.sleep(5)
    print(TASK_QUEUE.task_status())
    TASK_QUEUE.shutdown()
